package com.asiainfo.dacp.common.datasink;

import org.apache.flink.core.fs.Path;
import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.apache.flink.orc.writer.OrcBulkWriterFactory;
import org.apache.flink.streaming.api.functions.sink.filesystem.BucketAssigner;
import org.apache.flink.streaming.api.functions.sink.filesystem.OutputFileConfig;
import org.apache.flink.streaming.api.functions.sink.filesystem.StreamingFileSink;
import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.SimpleVersionedStringSerializer;
import org.apache.flink.streaming.api.functions.sink.filesystem.rollingpolicies.OnCheckpointRollingPolicy;

import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.UUID;

public class HdfsOrcSinkFunctionOpt<T> {

    /**
     * 根据schema和路径返回hdfsorcsink
     * @param path
     * @param schema
     * @return
     */
    public  StreamingFileSink<T> createOrcSinkFunction(String path, String schema){
        //根据schema生成向量构造器
        OrcBulkWriterFactory<T> factory = new OrcBulkWriterFactory<>(new PoJoGeneralOrcVector<>(schema));
        SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd");
        return   StreamingFileSink
                //指定存储路径和向量生成对象
                .forBulkFormat(new Path(path),factory)
                //分区策略
                .withBucketAssigner(new BucketAssigner<T, String>() {
                    @Override
                    public String getBucketId(T data, Context context) {
                        return "dt="+sdf.format(new Date());
                    }

                    @Override
                    public SimpleVersionedSerializer<String> getSerializer() {
                        return null;
                    }
                })
                //文件生成策略，前缀后缀等
                .withOutputFileConfig(OutputFileConfig.builder()
                        .withPartPrefix(UUID.randomUUID().toString())
                        .withPartSuffix(".dat")
                        .build())
                //滚动策略等
                .withRollingPolicy(OnCheckpointRollingPolicy.build())
                .build();
    }
}
